#!/usr/bin/env python3
"""
超时监控和问题诊断工具
专门解决系统卡住、长时间无响应的问题
"""

import asyncio
import time
import threading
import traceback
import psutil
import os
from datetime import datetime, timedelta
from typing import Dict, Any, Optional, Callable, List
from dataclasses import dataclass
from enum import Enum
import json

from logger_config import hr_logger

class TimeoutType(Enum):
    """超时类型"""
    MODEL_INVOKE = "MODEL_INVOKE"      # 大模型调用超时
    TOOL_EXECUTION = "TOOL_EXECUTION"  # 工具执行超时
    AGENT_PROCESSING = "AGENT_PROCESSING"  # 智能体处理超时
    TOTAL_REQUEST = "TOTAL_REQUEST"    # 总请求超时
    DATABASE_QUERY = "DATABASE_QUERY"  # 数据库查询超时

@dataclass
class TimeoutConfig:
    """超时配置"""
    model_timeout: float = 30.0          # 大模型调用超时(秒)
    tool_timeout: float = 10.0           # 工具执行超时(秒) 
    agent_timeout: float = 60.0          # 智能体处理超时(秒)
    total_timeout: float = 120.0         # 总请求超时(秒)
    database_timeout: float = 5.0       # 数据库超时(秒)

@dataclass
class TimeoutEvent:
    """超时事件记录"""
    event_id: str
    timeout_type: TimeoutType
    start_time: datetime
    timeout_threshold: float
    context: Dict[str, Any]
    stack_trace: Optional[str] = None
    resolved: bool = False
    resolution_time: Optional[datetime] = None

class TimeoutMonitor:
    """超时监控器"""
    
    def __init__(self, config: TimeoutConfig = None):
        self.config = config or TimeoutConfig()
        self.active_operations: Dict[str, TimeoutEvent] = {}
        self.timeout_history: List[TimeoutEvent] = []
        self.monitoring = True
        self.monitor_thread = None
        self._start_monitoring()
    
    def _start_monitoring(self):
        """启动监控线程"""
        def monitor_loop():
            while self.monitoring:
                try:
                    self._check_timeouts()
                    time.sleep(1)  # 每秒检查一次
                except Exception as e:
                    hr_logger.error("超时监控线程异常", exception=e)
        
        self.monitor_thread = threading.Thread(target=monitor_loop, daemon=True)
        self.monitor_thread.start()
        hr_logger.info("超时监控线程已启动")
    
    def _check_timeouts(self):
        """检查超时的操作"""
        current_time = datetime.now()
        
        for op_id, event in list(self.active_operations.items()):
            elapsed = (current_time - event.start_time).total_seconds()
            
            if elapsed > event.timeout_threshold:
                # 发现超时
                self._handle_timeout(op_id, event, elapsed)
    
    def _handle_timeout(self, op_id: str, event: TimeoutEvent, elapsed_time: float):
        """处理超时事件"""
        hr_logger.error(f"检测到超时事件", extra_data={
            "operation_id": op_id,
            "timeout_type": event.timeout_type.value,
            "elapsed_time": elapsed_time,
            "threshold": event.timeout_threshold,
            "context": event.context
        })
        
        # 获取当前线程堆栈
        for thread_id, frame in dict(threading._current_frames()).items():
            if thread_id == threading.current_thread().ident:
                continue
            stack = traceback.extract_stack(frame)
            event.stack_trace = ''.join(traceback.format_list(stack))
        
        # 移到历史记录
        self.timeout_history.append(event)
        del self.active_operations[op_id]
        
        # 记录系统状态
        self._log_system_state(event)
        
        # 尝试恢复操作
        self._attempt_recovery(event)
    
    def _log_system_state(self, event: TimeoutEvent):
        """记录系统状态用于诊断"""
        try:
            # CPU和内存使用情况
            cpu_percent = psutil.cpu_percent()
            memory = psutil.virtual_memory()
            
            # 进程信息
            process = psutil.Process(os.getpid())
            process_info = {
                "pid": process.pid,
                "cpu_percent": process.cpu_percent(),
                "memory_percent": process.memory_percent(),
                "num_threads": process.num_threads(),
                "open_files": len(process.open_files()),
                "connections": len(process.connections())
            }
            
            system_state = {
                "timestamp": datetime.now().isoformat(),
                "cpu_percent": cpu_percent,
                "memory_percent": memory.percent,
                "memory_available_gb": memory.available / (1024**3),
                "process_info": process_info,
                "timeout_event": {
                    "type": event.timeout_type.value,
                    "context": event.context
                }
            }
            
            hr_logger.error("系统状态快照", extra_data=system_state)
            
        except Exception as e:
            hr_logger.error("获取系统状态失败", exception=e)
    
    def _attempt_recovery(self, event: TimeoutEvent):
        """尝试恢复操作"""
        recovery_actions = {
            TimeoutType.MODEL_INVOKE: self._recover_model_invoke,
            TimeoutType.TOOL_EXECUTION: self._recover_tool_execution,
            TimeoutType.AGENT_PROCESSING: self._recover_agent_processing,
            TimeoutType.TOTAL_REQUEST: self._recover_total_request
        }
        
        recovery_action = recovery_actions.get(event.timeout_type)
        if recovery_action:
            try:
                recovery_action(event)
            except Exception as e:
                hr_logger.error(f"恢复操作失败: {event.timeout_type.value}", exception=e)
    
    def _recover_model_invoke(self, event: TimeoutEvent):
        """恢复大模型调用超时"""
        hr_logger.warning("大模型调用超时，建议检查:", extra_data={
            "suggestions": [
                "检查网络连接",
                "检查API密钥是否有效",
                "检查请求是否过大",
                "考虑重试或使用备用模型"
            ],
            "context": event.context
        })
    
    def _recover_tool_execution(self, event: TimeoutEvent):
        """恢复工具执行超时"""
        hr_logger.warning("工具执行超时，建议检查:", extra_data={
            "suggestions": [
                "检查工具参数是否正确",
                "检查数据库连接",
                "检查文件系统权限",
                "检查工具依赖"
            ],
            "context": event.context
        })
    
    def _recover_agent_processing(self, event: TimeoutEvent):
        """恢复智能体处理超时"""
        hr_logger.warning("智能体处理超时，建议检查:", extra_data={
            "suggestions": [
                "检查智能体逻辑是否有死循环",
                "检查状态管理是否正确",
                "检查消息传递链路",
                "考虑简化处理逻辑"
            ],
            "context": event.context
        })
    
    def _recover_total_request(self, event: TimeoutEvent):
        """恢复总请求超时"""
        hr_logger.warning("总请求超时，建议检查:", extra_data={
            "suggestions": [
                "检查请求复杂度",
                "检查系统负载",
                "考虑分解请求",
                "增加超时时间"
            ],
            "context": event.context
        })
    
    def start_operation(self, operation_id: str, timeout_type: TimeoutType, 
                       context: Dict[str, Any] = None, 
                       custom_timeout: float = None) -> str:
        """开始监控操作"""
        timeout_map = {
            TimeoutType.MODEL_INVOKE: self.config.model_timeout,
            TimeoutType.TOOL_EXECUTION: self.config.tool_timeout,
            TimeoutType.AGENT_PROCESSING: self.config.agent_timeout,
            TimeoutType.TOTAL_REQUEST: self.config.total_timeout,
            TimeoutType.DATABASE_QUERY: self.config.database_timeout
        }
        
        timeout_threshold = custom_timeout or timeout_map[timeout_type]
        
        event = TimeoutEvent(
            event_id=operation_id,
            timeout_type=timeout_type,
            start_time=datetime.now(),
            timeout_threshold=timeout_threshold,
            context=context or {}
        )
        
        self.active_operations[operation_id] = event
        
        hr_logger.debug(f"开始监控操作: {timeout_type.value}", extra_data={
            "operation_id": operation_id,
            "timeout_threshold": timeout_threshold,
            "context": context
        })
        
        return operation_id
    
    def end_operation(self, operation_id: str, success: bool = True):
        """结束监控操作"""
        if operation_id in self.active_operations:
            event = self.active_operations[operation_id]
            event.resolved = True
            event.resolution_time = datetime.now()
            
            elapsed = (event.resolution_time - event.start_time).total_seconds()
            
            if success:
                hr_logger.debug(f"操作完成: {event.timeout_type.value}", extra_data={
                    "operation_id": operation_id,
                    "elapsed_time": elapsed,
                    "threshold": event.timeout_threshold
                })
            else:
                hr_logger.warning(f"操作失败: {event.timeout_type.value}", extra_data={
                    "operation_id": operation_id,
                    "elapsed_time": elapsed
                })
            
            del self.active_operations[operation_id]
    
    def get_status(self) -> Dict[str, Any]:
        """获取监控状态"""
        current_time = datetime.now()
        
        active_ops = []
        for op_id, event in self.active_operations.items():
            elapsed = (current_time - event.start_time).total_seconds()
            active_ops.append({
                "operation_id": op_id,
                "type": event.timeout_type.value,
                "elapsed_time": elapsed,
                "threshold": event.timeout_threshold,
                "progress_percent": min(100, (elapsed / event.timeout_threshold) * 100)
            })
        
        recent_timeouts = [
            {
                "type": event.timeout_type.value,
                "start_time": event.start_time.isoformat(),
                "context": event.context
            }
            for event in self.timeout_history[-10:]  # 最近10个超时
        ]
        
        return {
            "monitoring_active": self.monitoring,
            "active_operations_count": len(self.active_operations),
            "active_operations": active_ops,
            "total_timeouts": len(self.timeout_history),
            "recent_timeouts": recent_timeouts,
            "config": {
                "model_timeout": self.config.model_timeout,
                "tool_timeout": self.config.tool_timeout,
                "agent_timeout": self.config.agent_timeout,
                "total_timeout": self.config.total_timeout
            }
        }
    
    def stop_monitoring(self):
        """停止监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=2)
        hr_logger.info("超时监控已停止")

class TimeoutContext:
    """超时上下文管理器"""
    
    def __init__(self, monitor: TimeoutMonitor, timeout_type: TimeoutType, 
                 operation_name: str, context: Dict[str, Any] = None,
                 custom_timeout: float = None):
        self.monitor = monitor
        self.timeout_type = timeout_type
        self.operation_name = operation_name
        self.context = context or {}
        self.custom_timeout = custom_timeout
        self.operation_id = None
    
    def __enter__(self):
        self.operation_id = f"{self.operation_name}_{int(time.time() * 1000)}"
        self.monitor.start_operation(
            self.operation_id, 
            self.timeout_type, 
            self.context, 
            self.custom_timeout
        )
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.operation_id:
            success = exc_type is None
            self.monitor.end_operation(self.operation_id, success)

# 全局超时监控实例
timeout_monitor = TimeoutMonitor()

# 装饰器
def monitor_timeout(timeout_type: TimeoutType, custom_timeout: float = None):
    """超时监控装饰器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            operation_name = f"{func.__module__}.{func.__name__}"
            context = {
                "function": func.__name__,
                "module": func.__module__,
                "args_count": len(args),
                "kwargs_keys": list(kwargs.keys())
            }
            
            with TimeoutContext(timeout_monitor, timeout_type, operation_name, context, custom_timeout):
                return func(*args, **kwargs)
        
        return wrapper
    return decorator

def async_monitor_timeout(timeout_type: TimeoutType, custom_timeout: float = None):
    """异步超时监控装饰器"""
    def decorator(func):
        async def wrapper(*args, **kwargs):
            operation_name = f"{func.__module__}.{func.__name__}"
            context = {
                "function": func.__name__,
                "module": func.__module__,
                "args_count": len(args),
                "kwargs_keys": list(kwargs.keys())
            }
            
            with TimeoutContext(timeout_monitor, timeout_type, operation_name, context, custom_timeout):
                return await func(*args, **kwargs)
        
        return wrapper
    return decorator

if __name__ == "__main__":
    # 测试超时监控
    import time
    
    monitor = TimeoutMonitor(TimeoutConfig(model_timeout=2.0))
    
    # 测试正常操作
    with TimeoutContext(monitor, TimeoutType.MODEL_INVOKE, "test_normal", {"test": "normal"}):
        time.sleep(1)
        print("正常操作完成")
    
    # 测试超时操作
    try:
        with TimeoutContext(monitor, TimeoutType.MODEL_INVOKE, "test_timeout", {"test": "timeout"}):
            time.sleep(3)  # 超过2秒超时
            print("这不应该打印")
    except:
        pass
    
    # 等待监控检测
    time.sleep(3)
    
    # 查看状态
    status = monitor.get_status()
    print(json.dumps(status, indent=2, ensure_ascii=False))
    
    monitor.stop_monitoring()












